"""
TODO: 多线程resize
"""
import argparse
import os
from PIL import Image

import pdb

COCO_PATH = "/home/dandelight/coco2017/coco2017"

def resize_image(image, size):
    """Resize an image to the given size."""
    return image.resize(size, Image.ANTIALIAS)

def resize_images(image_dir, output_dir, size):
    """Resize the images in `image_dir` and save to `output_dir`"""
    print("output dir: {}".format(output_dir))

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    images = os.listdir(image_dir)
    num_images = len(images)
    for i, image in enumerate(images):
        with open(os.path.join(image_dir, image), 'r+b') as f:
            with Image.open(f) as img:
                img = resize_image(img, size)
                # pdb.set_trace()
                img.save(os.path.join(output_dir, image))#, format=image.format)
            if (i+1) % 1000 == 0:
                print("[{}/{}] images resized.".format(i+1, num_images))
    print("DONE! Finished with the {} images.".format(num_images))

def main(args):
    image_dir = args.image_dir
    output_dir = args.output_dir
    image_size = [args.image_size, args.image_size]
    resize_images(image_dir, output_dir, image_size)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--image_dir', type=str, default='{}/train2017/'.format(COCO_PATH),
                        help='directory for train images')
    parser.add_argument('--output_dir', type=str, default='{}/resized2017/'.format(COCO_PATH),
                        help='directory for saving resized images')
    parser.add_argument('--image_size', type=int, default=256,
                        help='size for image after processing')
    args = parser.parse_args()
    main(args)
